const std = @import("std");
const os = @import("os.zig");
const PGSIZE = os.PGSIZE;
const KSTACK = os.KSTACK;
const log = std.log.scoped(.proc);

pub const Cpu = @import("proc/Cpu.zig");
pub const Proc = @import("proc/Proc.zig");

pub const pid_t = usize;
pub const max_pid = std.math.maxInt(std.meta.Int(.signed, @bitSizeOf(pid_t)));

/// Actual number of CPUs detected at runtime (may be less than NCPU)
pub var ncpus: usize = os.NCPU;
pub var all_procs: [os.NPROC]Proc = @splat(.{});
pub var init_proc: *Proc = &all_procs[0];
pub var kthreadd: *Proc = &all_procs[1];

/// Initializes the process table and creates essential system processes.
pub fn init() void {
    for (&all_procs, 1..) |*p, i| {
        p.state = .unused;
        p.kstack = os.mem.VirtAddr.TRAMPOLINE - (i * (KSTACK + 1) - KSTACK) * PGSIZE;
    }

    // Create initial system process with root directory
    errdefer @panic("bootstrap proc");
    init_proc = try alloc(.uproc);
    init_proc.ctx = .{ .uproc = try .create() };

    // Initialize kernel thread daemon - manages all worker threads
    // Becomes its own parent to enable waiting on children via waitAll()
    kthreadd = kthreadd.spawn("kthreadd", .{});
}

test "bootstrap processes" {
    try std.testing.expect(init_proc == &all_procs[0]);
    try std.testing.expect(kthreadd == &all_procs[1]);

    try std.testing.expect(init_proc.pid == 1);
    try std.testing.expect(kthreadd.pid == 2);
}

var pid_lock: os.Lock.Spin = .init("nextpid");
var nextpid: pid_t = 1;

fn allocpid() pid_t {
    pid_lock.acquire();
    defer pid_lock.release();

    const pid = nextpid;
    if (pid >= max_pid) @panic("pid overflow");

    nextpid += 1;
    return pid;
}

/// Allocate and initialize a process from the process table.
/// If successful, returns a process with its lock released and state set to used.
pub fn alloc(tag: Proc.Type) !*Proc {
    const p: *Proc = for (&all_procs) |*p| {
        p.lock.acquire();

        if (p.state == .unused) {
            break p;
        } else {
            p.lock.release();
        }
    } else {
        return error.SystemResources;
    };
    defer p.lock.release();

    p.pid = allocpid();
    p.cwd = os.fs.path.getinode("/") catch unreachable;
    p.context.sp = p.kstack;
    p.context.ra = if (tag == .kthread) &ktentry else &forkret;

    p.state = .used;
    return p;
}

/// Kernel thread entry point - handles thread lifecycle and cleanup.
fn ktentry() callconv(.c) noreturn {
    const p = current().?;
    const kthread = &p.ctx.kthread;

    // Still holding p.lock from scheduler.
    p.lock.release();

    // Wait for all thread dependencies
    p.waitAll();

    // Execute the actual kernel thread function
    kthread.func(kthread.arg);

    // Jump into the scheduler, never to return.
    p.exit(0);
    @panic("kthread exit");
}

/// A fork child's very first scheduling by scheduler() will swtch to forkret.
fn forkret() callconv(.c) noreturn {
    // Still holding p.lock from scheduler.
    current().?.lock.release();

    // return to user space, mimicing usertrap()'s return.
    os.trap.usertrapret();
}

/// Return the current process, or zero if none.
pub fn current() ?*Proc {
    const this = Cpu.push();
    defer this.pop();

    return this.process;
}

/// Per-CPU process scheduler. Each CPU calls scheduler() after setting itself up.
/// Scheduler never returns. It loops, doing:
///  - choose a process to run.
///  - swtch to start running that process.
///  - eventually that process transfers control via swtch back to the scheduler.
pub fn scheduler() noreturn {
    // This executes in the CPU's dedicated scheduler context
    // The CPU core remains constant throughout scheduler execution
    const cpu = Cpu.this();
    cpu.process = null;

    while (true) {
        // The most recent process to run may have had interrupts turned off; enable them to avoid a
        // deadlock if all processes are waiting.
        os.arch.intr.enable();
        // Then turn them back off to avoid a possible race between an interrupt and wfi.
        _ = os.arch.intr.disable();

        var found: bool = false;
        for (&all_procs) |*p| {
            p.lock.acquire();
            defer p.lock.release();

            if (p.state == .runnable) {
                // Switch to chosen process. It is the process's job to release its lock and then
                // reacquire it before jumping back to us.
                p.state = .running;
                cpu.process = p;
                cpu.context.swtch(&p.context);

                // Process is done running for now. It should have changed its p.state before coming back.
                cpu.process = null;
                found = true;
            }
        }
        if (!found) {
            // nothing to run; stop running on this core until an interrupt.
            asm volatile ("wfi");
        }
    }
}

/// Wake up all processes sleeping on channel `chan`. Caller should hold the condition lock.
pub fn wakeup(chan: *anyopaque) void {
    for (&all_procs) |*p| {
        if (p == current()) continue;

        p.lock.acquire();
        defer p.lock.release();

        if (p.state == .sleeping and p.chan == chan) {
            p.state = .runnable;
        }
    }
}

/// Kill the process with the given pid.
/// The victim won't exit until it tries to return to user space (see usertrap()).
pub fn kill(pid: pid_t) !void {
    for (&all_procs) |*p| {
        p.lock.acquire();
        defer p.lock.release();

        if (p.pid == pid) {
            p.killed = true;
            if (p.state == .sleeping) {
                p.state = .runnable;
            }
            return;
        }
    }

    return error.ProcessNotFound;
}

/// Print a process listing to console. For debugging.
/// Runs when user types ^P on console. No lock to avoid wedging a stuck machine further.
pub fn dump() void {
    // We iterate over array by reference because this will take less memory on stack.
    for (&all_procs) |*p| {
        if (p.state == .unused) continue;
        log.info(
            "{} {t} {s}",
            .{ p.pid, p.state, p.name.get() },
        );
    }
}

/// Starts the main system process after all modules are initialized.
pub fn start(_: ?*const anyopaque) void {
    // create device nodes
    const p = current().?;
    const sys = os.syscall.sysfile;
    sys.mkdir(p, "/dev") catch {};
    sys.mknod(p, "/dev/null", os.NULL, 0) catch {};
    sys.mknod(p, "/dev/console", os.CONSOLE, 0) catch {};

    // run tests instead of starting init
    if (os.is_test) @import("root").tests();

    // We can invoke exec() now that file system is initialized.
    const argc = init_proc.exec("/bin/init", &.{"init"}) catch |err| @panic(@errorName(err));
    init_proc.ctx.uproc.trapframe.reg.a0 = argc; // put argc into a0

    // Open /dev/console, for stdin/stdout/stderr, this should never fail
    const stdin = sys.open(init_proc, "/dev/console", .{ .access = .RDWR }) catch unreachable;
    const stdout = sys.dup(init_proc, stdin) catch unreachable;
    const stderr = sys.dup(init_proc, stdout) catch unreachable;
    std.debug.assert(stderr == 2);

    init_proc.run();
}

pub fn shutdown(code: u16) noreturn {
    @branchHint(.cold);
    os.driver.uart.freeze = true;
    os.driver.syscon.shutdown(if (code == 0) .pass else .fail, code);
}
